import onnxruntime as rt
import numpy as np
import cv2
import os
import time

# 获取图片路径
current_dir = os.path.dirname(os.path.abspath(__file__))
images_path = os.path.join(current_dir, "../data")
if not os.path.exists(images_path):
        raise Exception("the directory is not exist")
all_path = []
image_name = []
for path in os.listdir(images_path):
    if path != '.keep':
        total_path = os.path.join(images_path, path)
        all_path.append(total_path)
if len(all_path) == 0:
    raise Exception("the directory is empty, please download image")
    
# 推理所有图片
for input_path in all_path:
    input_path = os.path.abspath(input_path)
    
    # 预处理
    image = cv2.imread(input_path)
    # 图片缩放到28*28
    resize_image = cv2.resize(image, (28, 28))
    # 转换灰度图
    resize_image = cv2.cvtColor(resize_image, cv2.COLOR_BGR2GRAY)
    # 归一化
    input = (resize_image / 255).astype(np.float32)
    # 设置NCHW格式
    input = input.reshape(1, 1, 28, 28)
    
    # onnx推理配置项
    options = rt.SessionOptions()
    options.intra_op_num_threads = 1
    options.inter_op_num_threads = 1
    
    # 加载onnx模型和配置
    sess = rt.InferenceSession("../mnist.onnx",options)
    input_name = "input"
    output_name = "output"
    start_time = time.time()
    # 推理
    output = sess.run([output_name], {input_name: input})
    end_time = time.time()
    
    # 后处理
    output = np.array(output).flatten()
    # softmax
    vals = np.exp(output) / np.sum(np.exp(output))
    result = "[image_path:" + input_path + "] [inferssession_time:" + f"{int(1/(end_time - start_time))}" + " pictures/s]" + " [output:" + "\033[1;31;40m" + f"{np.argmax(vals)}" + "\033[0m]"
    print(result)


